import os
import cv2
import utils
import random
import numpy as np
from tqdm import tqdm

SrcDir = r"I:\ImgOutput\temp"
ReplaceLabelDir = r"I:\Temp5"
DstDir = r"I:\Temp4"
CoverName = ["0", "5"]
CoverNum = [5, 5]
ReplaceLabelName = ["2", "3", "4", "6", "7"]

RewriteKey = False

LabelNames = [[]] * len(ReplaceLabelName)
for root, dirs, files in os.walk(ReplaceLabelDir):
    if root[-1] in ReplaceLabelName:
        random.shuffle(files)
        Index = ReplaceLabelName.index(root[-1])
        LabelNames[Index] = files

for root, dirs, files in os.walk(SrcDir):
    random.shuffle(files)
    for file in tqdm(files):
        if file[-1] == 'l':
            Infos = utils.ReadXml(root + "\\" + file)
            NewInfos = [DstDir + "\\" + file]
            for Info in Infos:
                if Info[-1] in CoverName and CoverNum[CoverName.index(Info[-1])] > 0:
                    Img = cv2.imread(root + "\\" + file[:-3] + "jpg")
                    CoverNum[CoverName.index(Info[-1])] -= 1
                    RewriteKey = True
                    if len(ReplaceLabelName) == 0:
                        Img[Info[1]:Info[3], Info[0]:Info[2]] = np.ones(
                            (Info[3] - Info[1], Info[2] - Info[0], 3)) * np.random.randint(0, 255)
                    else:
                        NumberDir = random.randint(0, len(ReplaceLabelName) - 1)
                        NumberIndex = random.randint(0, len(LabelNames[NumberDir]) - 1)
                        ImgNumber = cv2.imread(ReplaceLabelDir + "\\" + ReplaceLabelName[NumberDir] + "\\" + LabelNames[NumberDir][NumberIndex])
                        ImgLabel = cv2.resize(ImgNumber, (Info[2] - Info[0], Info[3] - Info[1]))
                        Img[Info[1]:Info[3], Info[0]:Info[2]] = utils.BrightnessAverage(Img, ImgLabel)

                        cv2.imwrite(DstDir + "\\" + file[:-4] + "A.jpg", Img)
                        Img[Info[1]:Info[3], Info[0]:Info[2]] = ImgLabel
                        Info[-1] = ReplaceLabelName[NumberDir]
                        NewInfos.append(Info)
                else:
                    NewInfos.append(Info)
            if RewriteKey:
                RewriteKey = False
                utils.WriteXml(NewInfos, Img.shape[1], Img.shape[0])
                cv2.imwrite(DstDir + "\\" + file[:-3] + "jpg", Img)
